﻿using System.Collections.Generic;
using ProjectDanmaku.Core;
using ProjectDanmaku.Util;
using UnityEngine;
using UnityEngine.Rendering;
using UnityEngine.Rendering.Universal;

namespace ProjectDanmaku.Entity.Rendering
{
    internal struct ShaderPropertiesPerEntity
    {
        public Vector4 position;
        public float rotation;
        public int color;

        public static int Size()
        {
            return sizeof(float) * 6;
        }
    }
    
    public class EntityRenderFeature : ScriptableRendererFeature
    {
        private static readonly int ShaderPerEntityProperties = Shader.PropertyToID("_ShaderPerEntityProperties");
        private LayerMask danmakuLayerMask;

        public static Color? EntityColorOverride = null;
        
        private class EntityRenderPass : ScriptableRenderPass
        {
            private static Bounds _renderBounds = new Bounds(Vector3.zero, new Vector3(100, 100, 100));
            private Stack<ComputeBuffer> _computeBufferStack = new ();
            private uint[] _argsArray = new uint[5] { 0, 0, 0, 0, 0 };

            private int _onceRenderedEntityCount;

            public override void Execute(ScriptableRenderContext context, ref RenderingData renderingData)
            {
                ReleaseAllComputeBufferInStack();
                
                var cmd = CommandBufferPool.Get("Entity SubLayer Batch Renderer");
                var subLayers = FindObjectsByType<EntitySubLayer>(FindObjectsSortMode.None);

                _onceRenderedEntityCount = 0;
                
                foreach (var subLayer in subLayers)
                {
                    if (!subLayer.enabled)
                        continue;
                    
                    FillCommandForEntitySubLayer(cmd, subLayer);
                }
                
                context.ExecuteCommandBuffer(cmd);

                GlobalSharedData.RenderedEntityCount = _onceRenderedEntityCount;
            }

            private void ReleaseAllComputeBufferInStack()
            {
                while (_computeBufferStack.Count > 0)
                {
                    _computeBufferStack.Pop().Release();
                }
            }

            private int FillPerEntityComputeBuffer(
                Material material, List<EntityData> entities, bool doColorOverride, Color overrideColor)
            {
                var dataList = new List<ShaderPropertiesPerEntity>();

                foreach (var entity in entities)
                {
                    if (entity.Active)
                    {
                        dataList.Add(new ShaderPropertiesPerEntity()
                        {
                            position = entity.Position,
                            rotation = entity.RotateZ,
                            color = ColorPacker.PackColor(doColorOverride ? overrideColor : entity.Color)
                        });
                    }
                }

                if (dataList.Count == 0)
                    return 0;
                
                var computeBuffer = new ComputeBuffer(dataList.Count, ShaderPropertiesPerEntity.Size());
                computeBuffer.SetData(dataList);
                
                _computeBufferStack.Push(computeBuffer);;
                
                material.SetBuffer(ShaderPerEntityProperties, computeBuffer);
                
                return dataList.Count;
            }

            private void FillCommandForEntitySubLayer(CommandBuffer cmd, EntitySubLayer subLayer)
            {
                
                if (subLayer.Entities.Count <= 0)
                    return;
                
                var mesh = subLayer.GetSpriteMesh();
                var material = subLayer.GetMaterial();
                
                var colorOverride = Color.white;
                if (EntityColorOverride != null)
                    colorOverride = EntityColorOverride.Value;

                var instanceCount = FillPerEntityComputeBuffer(
                    material, subLayer.Entities, EntityColorOverride.HasValue, colorOverride);
                
                if (instanceCount == 0)
                    return;
                
                _onceRenderedEntityCount += instanceCount;

                _argsArray[0] = (uint)mesh.GetIndexCount(0);
                _argsArray[1] = (uint)instanceCount;
                _argsArray[2] = (uint)mesh.GetIndexStart(0);
                _argsArray[3] = (uint)mesh.GetBaseVertex(0);
                
                var argBuffer = new ComputeBuffer(
                    1, _argsArray.Length * sizeof(uint), ComputeBufferType.IndirectArguments);
                argBuffer.SetData(_argsArray);
                
                _computeBufferStack.Push(argBuffer);

                // cmd.BeginSample("Render EntitySubLayer " + subLayer.name);
                cmd.DrawMeshInstancedIndirect(
                    mesh, 0, material, 0, argBuffer);
                // cmd.EndSample("Render EntitySubLayer " + subLayer.name);
            }

            ~EntityRenderPass()
            {
                ReleaseAllComputeBufferInStack();
            }
        }

        private EntityRenderPass _renderPass;
        
        public override void Create()
        {
            danmakuLayerMask = LayerMask.NameToLayer("Danmaku");
            _renderPass = new EntityRenderPass()
            {
                renderPassEvent = RenderPassEvent.AfterRenderingTransparents
            };
        }

        public override void AddRenderPasses(ScriptableRenderer renderer, ref RenderingData renderingData)
        {
            if ((renderingData.cameraData.camera.cullingMask & danmakuLayerMask) != 0)
            {
                renderer.EnqueuePass(_renderPass);
            }
        }
    }
}